import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def _corrDisplayFactorBenchmark(corrData, factor, figsize=None):
    '''利用相关系数、p值构成的DataFrame画出条形图，x轴是因子'''
    fig = plt.figure(figsize=figsize)
    ax1 = fig.add_subplot(111)
    corrData.sort_values('r', ascending=False, inplace=True)
    corrData['xtick'] = range(len(corrData))
    sigPos = corrData[(corrData.r > 0) & (corrData.p < 0.01)]
    sigNeg = corrData[(corrData.r < 0) & (corrData.p < 0.01)]
    nonSig = corrData[~(corrData.p < 0.01)]
    ax1.bar(sigPos.xtick, sigPos.r, align='center', color='steelblue', label=u'显著为正')
    ax1.bar(nonSig.xtick, nonSig.r, align='center', color='powderblue', label=u'不显著')
    ax1.bar(sigNeg.xtick, sigNeg.r, align='center', color='deepskyblue', label=u'显著为负')
    plt.xticks(corrData.xtick, corrData.index)
    plt.ylim([-1, 1])
    plt.title(f"{factor}因子与风格因子的相关性")
    # plt.legend(loc='best')

    # plt.legend([u'显著为正', u'不显著', u'显著为负'], loc = 1, prop = legendFont)
    fig.autofmt_xdate(rotation=30, ha='right')
    plt.grid(axis='y')
    plt.show()


def ic_time_plot(corrData, figsize):
    fig = plt.figure(figsize=figsize)
    ax1 = fig.add_subplot(111)
    corrData['xtick'] = range(len(corrData))
    sigPos = corrData[(corrData.r > 0) & (corrData.p < 0.01)]
    sigNeg = corrData[(corrData.r < 0) & (corrData.p < 0.01)]
    nonSig = corrData[~(corrData.p < 0.01)]
    ax1.bar(sigPos.xtick, sigPos.r, align='center', color='steelblue', label=u'显著为正')
    ax1.bar(nonSig.xtick, nonSig.r, align='center', color='powderblue', label=u'不显著')
    ax1.bar(sigNeg.xtick, sigNeg.r, align='center', color='deepskyblue', label=u'显著为负')
    # plt.legend(loc='best')
    plt.xticks(corrData.xtick, corrData.index)
    plt.xlim(-1, len(corrData))
    fig.autofmt_xdate(rotation=70, ha='center')
    plt.grid(axis='y')
    plt.title(u'因子Rank IC值')
    plt.show()


def group_plot(factor, active_group_ret_dict, group_ret_dict, index_return, excess_return, figsize=None):
    fig = plt.figure(figsize=figsize)
    colormap = plt.cm.nipy_spectral
    colors = [colormap(i) for i in np.linspace(0, 1, len(group_ret_dict))]
    ax1 = fig.add_subplot(311)
    ax1.set_prop_cycle('color', colors)
    for group in active_group_ret_dict:
        ax1.plot(pd.to_datetime(active_group_ret_dict[group].index), active_group_ret_dict[group], label=group)
    ax1.set_title(f"{factor} 因子分层组合主动收益")
    ax1.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    ax1.grid()
    ax3 = fig.add_subplot(312)
    ax3.set_prop_cycle('color', colors)
    for group in group_ret_dict:
        time_index = pd.to_datetime(group_ret_dict[group].index)
        ax3.plot(time_index, (group_ret_dict[group].fillna(0) + 1).cumprod(), label=group)
    ax3.plot(time_index, (index_return + 1).cumprod(), '--', label='基准', )
    ax3.set_title(f"{factor}因子分层组合绝对收益")
    ax3.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    ax3.grid()

    ax2 = fig.add_subplot(313)
    ax2.plot(pd.to_datetime(excess_return.index), excess_return, label='多—空', color='steelblue')
    ax2.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    ax2.set_title("因子分层组合超额收益")
    ax2.grid()
    # CON_factor = quant_researcher.quant.project_tool.db_operator.db_conn.get_tk_factors_conn()
    # factor_CN = pd.read_sql("select factor_name_cn from stk_fac_info "
    #                         f"where factor_name_en = '{factor}'", CON_factor).iloc[0, 0]
    # plt.savefig(f'{factor_CN}.png')
    fig.show()


def industry_dist_plot(distDataFrame, factor):
    '''对于factorIndustryDist()给出的结果进行图形化，默认因子名称是self.factor'''
    fig = plt.figure(figsize=(16, 6))
    ax1 = fig.add_subplot(111)
    distDataFrame.sort_values('Q50', ascending=False).Q80.plot(kind='bar', ax=ax1, color='steelblue')
    distDataFrame.sort_values('Q50', ascending=False).Q50.plot(kind='bar', ax=ax1, color='powderblue')
    distDataFrame.sort_values('Q50', ascending=False).Q20.plot(kind='bar', ax=ax1, color='deepskyblue')

    plt.legend([u'80%分位点', u'中位数', u'20%分位点'])
    plt.xticks(rotation=70)
    plt.title(f'{factor}在各行业中分布情况')
    plt.grid(axis='y')
    plt.show()
